#include "point_Jacobi.hpp"
#include "scal_kernels.hpp"

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
PointJacobi<idx_t, data_t, setup_t, calc_t, dof>::PointJacobi(const calc_t wgt) : Solver<idx_t, setup_t, calc_t, dof>() {
    // this->weight = 0.8;
    this->weight = wgt;
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
PointJacobi<idx_t, data_t, setup_t, calc_t, dof>::~PointJacobi() {
    if (invD_separated ) { delete w_invDiag; w_invDiag = nullptr;}
    if (A_calc && own_A) { delete A_calc   ; A_calc    = nullptr; }
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointJacobi<idx_t, data_t, setup_t, calc_t, dof>::SetOperator(const Operator<idx_t, setup_t, dof> & op)
{
    assert(this->setup_called == false);
    this->setup_time = - wall_time();
    const par_CSRMatrix<idx_t, setup_t, setup_t, dof> & par_A = (const par_CSRMatrix<idx_t, setup_t, setup_t, dof> &) op;
    int my_pid; MPI_Comm_rank(par_A.comm, & my_pid);
    if (my_pid == 0) printf("Setup PJ #dof %d weight %.3f\n", dof, this->weight);
    
    separate_invD(par_A);

    if constexpr (sizeof(setup_t) == sizeof(calc_t) && sizeof(setup_t) == sizeof(data_t)) {
        own_A = false;
        A_calc = & par_A;
    }
    else {
        own_A = true;
        par_CSRMatrix<idx_t, data_t, calc_t, dof> * A_copy = new par_CSRMatrix<idx_t, data_t, calc_t, dof>
            (par_A.comm, par_A.rows_partition, par_A.cols_partition);
        // 拷贝数据，但不拷贝通信包
        A_copy->commpkg           = par_A.commpkg;
        A_copy->commpkg_transpose = par_A.commpkg_transpose;
        A_copy->col_map_offd      = par_A.col_map_offd;
        auto copy_trunc = [](const CSRMatrix<idx_t, setup_t, setup_t, dof> & src, CSRMatrix<idx_t, data_t, calc_t, dof> & dst) {
            dst.nrows = src.nrows;
            dst.ncols = src.ncols;
            dst.nnz   = src.nnz;
            dst.alloc_mem(false);

            #pragma omp parallel
            {
                #pragma omp for schedule(static) nowait
                for (idx_t i = 0; i <= dst.nrows; i++) {
                    dst.row_ptr[i] = src.row_ptr[i];
                }
                #pragma omp for schedule(static) nowait
                for (idx_t j = 0; j < dst.nnz; j++) {
                    dst.col_idx[j] = src.col_idx[j];
                }
                #pragma omp for schedule(static) nowait
                for (idx_t j = 0; j < dst.nnz * (long long) dof*dof; j++) {
                    dst.vals   [j] = src.vals   [j];
                }
            }
        };
        copy_trunc(par_A.diag, A_copy->diag);
        copy_trunc(par_A.offd, A_copy->offd);

        A_calc = A_copy;
    }
    this->setup_called = true;
    this->setup_time += wall_time();
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointJacobi<idx_t, data_t, setup_t, calc_t, dof>::truncate() {}

template<typename idx_t, typename data_t, typename calc_t, int dof>
static void elemwise_mul(const CSRMatrix<idx_t, data_t, calc_t, dof> & mat, seq_Vector<idx_t, calc_t, dof> & vec)
{
    assert(mat.nrows == vec.num);
    const data_t * mat_data = mat.vals;
    calc_t       * vec_data = vec.data;
    
    if constexpr (dof == 1) {
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < vec.num; i++) {
            vec_data[i] *= mat_data[i];
        }
    } else { static_assert(dof > 1);
        constexpr int e_size = dof*dof;
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < vec.num; i++) {
            matvec_mul<idx_t, data_t, calc_t, dof>(mat_data + i * e_size, vec_data + i * dof, 1.0);
        }
    }
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointJacobi<idx_t, data_t, setup_t, calc_t, dof>::Mult_ipl(const par_Vector<idx_t, calc_t, dof> & b, par_Vector<idx_t, calc_t, dof> & x) const
{
    if (this->zero_guess) {// x <--- w*D^{-1}*b
        vec_copy(b, x);// x <--- b
        elemwise_mul<idx_t, data_t, calc_t, dof>(*w_invDiag, *(x.local_vector));// x <--- w*D^{-1}*b
    }
    else {// x <--- x + w*D^{-1}*(b - A*x)
        par_Vector<idx_t, calc_t, dof> tmp(b);
        this->A_calc->Mult(x, tmp, false);
        vec_add(b, (calc_t) -1.0, tmp, tmp);// tmp <--- b - A*x
        elemwise_mul<idx_t, data_t, calc_t, dof>(*w_invDiag, *(tmp.local_vector));// tmp <--- w*D^{-1}*(b - A*x)
        vec_add(x, (calc_t) 1.0, tmp, x);// x <--- x + w*D^{-1}*(b - A*x)
    }
}

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
void PointJacobi<idx_t, data_t, setup_t, calc_t, dof>::separate_invD(const par_CSRMatrix<idx_t, setup_t, setup_t, dof> & par_A)
{
    assert(invD_separated == false);
    w_invDiag = new CSRMatrix<idx_t, data_t, calc_t, dof>;
    const CSRMatrix<idx_t, setup_t, setup_t, dof> & csr_A = par_A.diag;
    constexpr int e_size = dof*dof;

    w_invDiag->nrows = csr_A.nrows;
    w_invDiag->ncols = csr_A.ncols;
    w_invDiag->nnz   = w_invDiag->nrows;// 只存储主对角位置
    w_invDiag->vals  = new data_t [(long long)w_invDiag->nrows * e_size];

    for (idx_t i = 0; i < csr_A.nrows; i++) {
        
        for (auto p = csr_A.row_ptr[i]; p < csr_A.row_ptr[i + 1]; p++) {
            const idx_t j = csr_A.col_idx[p];
            if (j == i) {
                if constexpr (dof == 1) {
                    w_invDiag->vals[i] = this->weight / csr_A.vals[p];
                } else { static_assert(dof > 1);
                    setup_t A_copy[e_size];
                    data_t * dst_ptr = w_invDiag->vals + i * e_size;
                    for (int f = 0; f < e_size; f++) A_copy[f] = csr_A.vals[p * e_size + f];
                    // 计算该dof*dof的小矩阵的逆矩阵
                    if constexpr (sizeof(setup_t) == sizeof(data_t)) {// 不需要精度转换，直接写回
                        matinv_row<setup_t, dof>(A_copy, dst_ptr);
                        for (int f = 0; f < e_size; f++) dst_ptr[f] *= this->weight;
                    } else { static_assert(sizeof(data_t) < sizeof(setup_t));// 需要计算完后截断
                        setup_t invA_copy[e_size];
                        matinv_row<setup_t, dof>(A_copy, invA_copy);
                        for (int f = 0; f < e_size; f++) dst_ptr[f]  = this->weight * invA_copy[f];
                    }
                }
            }
        }
        
    }
    invD_separated = true;
}

template class PointJacobi<int, double, double, double, 1>;
template class PointJacobi<int, double, double, double, 2>;
template class PointJacobi<int, double, double, double, 3>;
template class PointJacobi<int, double, double, double, 4>;

template class PointJacobi<int, float , double, double, 1>;
template class PointJacobi<int, float , double, double, 2>;
template class PointJacobi<int, float , double, double, 3>;
template class PointJacobi<int, float , double, double, 4>;

template class PointJacobi<int, float , double, float , 1>;
template class PointJacobi<int, float , double, float , 2>;
template class PointJacobi<int, float , double, float , 3>;
template class PointJacobi<int, float , double, float , 4>;

template class PointJacobi<int, float , float , float , 1>;
template class PointJacobi<int, float , float , float , 2>;
template class PointJacobi<int, float , float , float , 3>;
template class PointJacobi<int, float , float , float , 4>;
#ifdef USE_FP16
template class PointJacobi<int, __fp16, double, float , 1>;
template class PointJacobi<int, __fp16, double, float , 2>;
template class PointJacobi<int, __fp16, double, float , 3>;
template class PointJacobi<int, __fp16, double, float , 4>;

template class PointJacobi<int, __fp16, float , float , 1>;
template class PointJacobi<int, __fp16, float , float , 2>;
template class PointJacobi<int, __fp16, float , float , 3>;
template class PointJacobi<int, __fp16, float , float , 4>;
#endif
